import torch
import os

# 检查 PyTorch 版本
print(torch.__version__)  # 1.7.1
# 查看cuda的版本
print(torch.version.cuda) # 10.2
# cudnn版本查询
print(torch.backends.cudnn.version()) #7605
# 查看显卡的型号
print(torch.cuda.get_device_name())  # GeForce GTX 1080

# 更新 PyTorch
# PyTorch 将被安装在 anaconda3/lib/python3.7/site-packages/torch/目录下。
# conda update pytorch torchvision -c pytorch

# 固定随机种子
torch.manual_seed(0)
torch.cuda.manual_seed(0)

# 指定程序运行在特定 GPU 卡上
# 在命令行指定环境变量
# CUDA_VISIBLE_DEVICES=0,1 python train.py

# 或在代码中指定
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

# 设置为 cuDNN benchmark 模式
# Benchmark 模式会提升计算速度，但是由于计算中有随机性，每次网络前馈结果略有差异。
torch.backends.cudnn.benchmark = True
# 如果想要避免这种结果波动，设置
torch.backends.cudnn.deterministic = True


# 清除 GPU 存储dedd
# 有时 Control-C 中止运行后 GPU 存储没有及时释放，需要手动清空。在 PyTorch 内部可以
torch.cuda.empty_cache()

# 或在命令行可以先使用 ps 找到程序的 PID，再使用 kill 结束该进程
# ps aux | grep pythonkill -9 [pid]

# ps aux | grep pythonkill -9 [pid]
# nvidia-smi --gpu-reset -i [gpu_id]
